import psycopg2
import json
from torch.utils.data import Dataset, DataLoader


def getTrainData(page=0, size=50000):
    # 连接到PostgreSQL数据库
    conn = psycopg2.connect(
        dbname="gserver", user="postgres", password="postgres", host="127.0.0.1"
    )
    cur = conn.cursor()
    # 执行查询获取bytea字段的数据
    cur.execute(
        # "SELECT ranges,train_data,test_data FROM stock_line_train_data"  # where code='0.002432' order by code"
        "SELECT ranges,train_data,test_data FROM stock_line_train_data order by id limit "
        + str(size)
        + " offset "
        + str(page * size)
    )
    rows = cur.fetchall()
    # ranges_bytea = cur.fetchone()[0]

    train_datas = []
    test_datas = []
    ranges = []
    try:
        # if ranges_bytea != any:
        #     ranges = json.loads(ranges_bytea)
        # else:
        #     ranges = ""

        for row in rows:
            ranges.append(json.loads(row[0]))

            train_bytea = bytes(row[1])
            train_bytea = train_bytea.decode("utf-8")
            train_datas.append(json.loads(train_bytea))
            test_bytea = bytes(row[2])
            test_bytea = test_bytea.decode("utf-8")
            test_datas.append(json.loads(test_bytea))

    except json.JSONDecodeError as e:
        print("Error decoding JSON data:", e)

    # 将Python列表转换为NumPy数组（如果需要）
    import numpy as np

    # 关闭游标和连接
    cur.close()
    conn.close()
    return ranges, train_datas, test_datas


def getTestData(page=0, size=50000):
    # 连接到PostgreSQL数据库
    conn = psycopg2.connect(
        dbname="gserver", user="postgres", password="postgres", host="127.0.0.1"
    )
    cur = conn.cursor()
    # 执行查询获取bytea字段的数据
    cur.execute(
        "SELECT code,end_date,ranges,train_data,test_data FROM stock_line_train_data order by code,end_date limit "
        + str(size)
        + " offset "
        + str(page * size)
    )
    rows = cur.fetchall()
    codes = []
    end_dates = []
    train_datas = []
    test_datas = []
    ranges = []
    try:
        for row in rows:
            codes.append(row[0])
            end_dates.append(row[1])
            ranges.append(json.loads(row[2]))
            train_bytea = bytes(row[3])
            train_bytea = train_bytea.decode("utf-8")
            train_datas.append(json.loads(train_bytea))
            test_bytea = bytes(row[4])
            test_bytea = test_bytea.decode("utf-8")
            test_datas.append(json.loads(test_bytea))

    except json.JSONDecodeError as e:
        print("Error decoding JSON data:", e)

    # 关闭游标和连接
    cur.close()
    conn.close()
    return codes, end_dates, ranges, train_datas, test_datas


def getPredData(page=0, size=50000):
    # 连接到PostgreSQL数据库
    conn = psycopg2.connect(
        dbname="gserver", user="postgres", password="postgres", host="127.0.0.1"
    )
    cur = conn.cursor()
    # 执行查询获取bytea字段的数据
    cur.execute(
        "SELECT code,end_date,train_data FROM stock_line_pred_data order by code,end_date limit "
        + str(size)
        + " offset "
        + str(page * size)
    )
    rows = cur.fetchall()
    codes = []
    end_dates = []
    train_datas = []
    try:
        for row in rows:
            codes.append(row[0])
            end_dates.append(row[1])
            train_bytea = bytes(row[2])
            train_bytea = train_bytea.decode("utf-8")
            train_datas.append(json.loads(train_bytea))

    except json.JSONDecodeError as e:
        print("Error decoding JSON data:", e)

    # 关闭游标和连接
    cur.close()
    conn.close()
    return codes, end_dates, train_datas


class Stock:
    def __init__(self, Code, Name):
        self.Code = Code
        self.Name = Name


def getStockInfos():
    # 连接到PostgreSQL数据库
    conn = psycopg2.connect(
        dbname="gserver2", user="postgres", password="postgres", host="127.0.0.1"
    )
    cur = conn.cursor()
    # 执行查询获取bytea字段的数据
    cur.execute("SELECT * from stock_infos")
    rows = cur.fetchall()
    stocks = {}
    try:
        for row in rows:
            stocks[row[0]] = row

    except json.JSONDecodeError as e:
        print("Error decoding JSON data:", e)

    # 关闭游标和连接
    cur.close()
    conn.close()
    return stocks


stocks = getStockInfos()
print(stocks["0.836149"])

# # # 测试
# ranges, train, test = getTrainData()
# # print(ranges, len(train[0]), len(test[0]))

# # dataset = TrainDataset(train, test)
# # data_loader = DataLoader(dataset, batch_size=10, shuffle=True)
# import torch
# from torch.utils.data import TensorDataset, DataLoader

# features = torch.tensor(train, dtype=torch.float32)
# # labels 是一个形状为 (num_samples,) 的一维数组，表示标签数据
# labels = torch.tensor(test, dtype=torch.float32)
# # 使用 TensorDataset 创建数据集
# dataset = TensorDataset(features, labels)
# # 可以选择使用 DataLoader 来批量加载数据集
# data_loader = DataLoader(dataset, batch_size=2, shuffle=True)

# # 现在可以迭代 data_loader 来获取批量数据
# for batch_features, batch_labels in data_loader:
#     # 在这里使用 batch_features 和 batch_labels
#     print(batch_features, batch_labels)
